package com.kwdb.sdk;

import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.exception.KWDBException;
import org.junit.Before;
import org.junit.Test;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 使用真实数据集进行性能测试
 */
public class RealDatasetTest {
    
    private KWDBClient client;
    private TestDatasetDownloader downloader;
    
    @Before
    public void setUp() throws KWDBException {
        ClientConfig config = new ClientConfig("localhost", 9090)
                .setBatchSize(1000)
                .setBatchTimeout(100)
                .setThreadPoolSize(20)
                .setEnableCompression(true);
        
        client = new KWDBClient(config);
        client.connect();
        
        downloader = new TestDatasetDownloader();
    }
    
    /**
     * 使用NOAA水数据库数据集进行测试
     */
    @Test
    public void testWithNOAADataset() throws IOException, KWDBException {
        System.out.println("=== NOAA水数据库数据集测试 ===");
        
        // 下载数据集
        String filepath = downloader.downloadNOAADataset();
        
        // 解析数据集（限制为10000行以节省时间）
        List<DataPoint> dataPoints = downloader.parseNOAADataset(filepath, 10000);
        
        if (dataPoints.isEmpty()) {
            System.out.println("数据集为空，跳过测试");
            return;
        }
        
        System.out.println("开始写入NOAA数据集，共 " + dataPoints.size() + " 个数据点");
        
        long startTime = System.currentTimeMillis();
        AtomicLong successCount = new AtomicLong(0);
        AtomicLong errorCount = new AtomicLong(0);
        
        // 分批写入
        int batchSize = 1000;
        for (int i = 0; i < dataPoints.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dataPoints.size());
            List<DataPoint> batch = dataPoints.subList(i, endIndex);
            
            try {
                client.writeBatch(batch);
                successCount.addAndGet(batch.size());
                
                if (successCount.get() % 5000 == 0) {
                    System.out.println("已写入: " + successCount.get() + " 个数据点");
                }
            } catch (Exception e) {
                errorCount.addAndGet(batch.size());
                System.err.println("批次写入失败: " + e.getMessage());
            }
        }
        
        long endTime = System.currentTimeMillis();
        long duration = endTime - startTime;
        double recordsPerSecond = (double) successCount.get() / duration * 1000;
        
        System.out.println("NOAA数据集测试完成:");
        System.out.println("成功写入: " + successCount.get() + " 个数据点");
        System.out.println("失败写入: " + errorCount.get() + " 个数据点");
        System.out.println("总耗时: " + duration + "ms");
        System.out.println("写入速度: " + String.format("%.2f", recordsPerSecond) + " 记录/秒");
        System.out.println("成功率: " + String.format("%.2f", (double) successCount.get() / dataPoints.size() * 100) + "%");
    }
    
    /**
     * 使用股票价格数据集进行测试
     */
    @Test
    public void testWithStockPriceDataset() throws IOException, KWDBException {
        System.out.println("=== 股票价格数据集测试 ===");
        
        // 下载数据集
        String filepath = downloader.downloadStockPriceDataset();
        
        // 解析CSV数据集
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 0); // 读取全部数据
        
        if (dataPoints.isEmpty()) {
            System.out.println("数据集为空，跳过测试");
            return;
        }
        
        System.out.println("开始写入股票价格数据集，共 " + dataPoints.size() + " 个数据点");
        
        long startTime = System.currentTimeMillis();
        AtomicLong successCount = new AtomicLong(0);
        
        // 分批写入
        int batchSize = 100;
        for (int i = 0; i < dataPoints.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dataPoints.size());
            List<DataPoint> batch = dataPoints.subList(i, endIndex);
            
            try {
                client.writeBatch(batch);
                successCount.addAndGet(batch.size());
            } catch (Exception e) {
                System.err.println("批次写入失败: " + e.getMessage());
            }
        }
        
        long endTime = System.currentTimeMillis();
        long duration = endTime - startTime;
        double recordsPerSecond = (double) successCount.get() / duration * 1000;
        
        System.out.println("股票价格数据集测试完成:");
        System.out.println("成功写入: " + successCount.get() + " 个数据点");
        System.out.println("总耗时: " + duration + "ms");
        System.out.println("写入速度: " + String.format("%.2f", recordsPerSecond) + " 记录/秒");
    }
    
    /**
     * 使用IoT传感器数据集进行测试
     */
    @Test
    public void testWithIoTSensorDataset() throws IOException, KWDBException {
        System.out.println("=== IoT传感器数据集测试 ===");
        
        // 下载数据集
        String filepath = downloader.downloadIoTSensorDataset();
        
        // 如果是压缩文件，先解压
        if (filepath.endsWith(".zip")) {
            System.out.println("检测到ZIP文件，需要解压处理");
            // 这里可以添加ZIP解压逻辑
            System.out.println("跳过ZIP文件测试");
            return;
        }
        
        // 解析数据集
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 5000); // 限制5000行
        
        if (dataPoints.isEmpty()) {
            System.out.println("数据集为空，跳过测试");
            return;
        }
        
        System.out.println("开始写入IoT传感器数据集，共 " + dataPoints.size() + " 个数据点");
        
        long startTime = System.currentTimeMillis();
        AtomicLong successCount = new AtomicLong(0);
        
        // 分批写入
        int batchSize = 500;
        for (int i = 0; i < dataPoints.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dataPoints.size());
            List<DataPoint> batch = dataPoints.subList(i, endIndex);
            
            try {
                client.writeBatch(batch);
                successCount.addAndGet(batch.size());
            } catch (Exception e) {
                System.err.println("批次写入失败: " + e.getMessage());
            }
        }
        
        long endTime = System.currentTimeMillis();
        long duration = endTime - startTime;
        double recordsPerSecond = (double) successCount.get() / duration * 1000;
        
        System.out.println("IoT传感器数据集测试完成:");
        System.out.println("成功写入: " + successCount.get() + " 个数据点");
        System.out.println("总耗时: " + duration + "ms");
        System.out.println("写入速度: " + String.format("%.2f", recordsPerSecond) + " 记录/秒");
    }
    
    /**
     * 使用网络监控数据集进行测试
     */
    @Test
    public void testWithNetworkMonitoringDataset() throws IOException, KWDBException {
        System.out.println("=== 网络监控数据集测试 ===");
        
        // 下载数据集
        String filepath = downloader.downloadNetworkMonitoringDataset();
        
        // 如果是GZ文件，先解压
        if (filepath.endsWith(".gz")) {
            System.out.println("检测到GZ文件，正在解压...");
            filepath = downloader.extractGZFile(filepath);
        }
        
        // 解析数据集（限制为10000行）
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 10000);
        
        if (dataPoints.isEmpty()) {
            System.out.println("数据集为空，跳过测试");
            return;
        }
        
        System.out.println("开始写入网络监控数据集，共 " + dataPoints.size() + " 个数据点");
        
        long startTime = System.currentTimeMillis();
        AtomicLong successCount = new AtomicLong(0);
        
        // 分批写入
        int batchSize = 1000;
        for (int i = 0; i < dataPoints.size(); i += batchSize) {
            int endIndex = Math.min(i + batchSize, dataPoints.size());
            List<DataPoint> batch = dataPoints.subList(i, endIndex);
            
            try {
                client.writeBatch(batch);
                successCount.addAndGet(batch.size());
            } catch (Exception e) {
                System.err.println("批次写入失败: " + e.getMessage());
            }
        }
        
        long endTime = System.currentTimeMillis();
        long duration = endTime - startTime;
        double recordsPerSecond = (double) successCount.get() / duration * 1000;
        
        System.out.println("网络监控数据集测试完成:");
        System.out.println("成功写入: " + successCount.get() + " 个数据点");
        System.out.println("总耗时: " + duration + "ms");
        System.out.println("写入速度: " + String.format("%.2f", recordsPerSecond) + " 记录/秒");
    }
    
    /**
     * 综合数据集性能测试
     */
    @Test
    public void testWithMultipleDatasets() throws IOException, KWDBException {
        System.out.println("=== 综合数据集性能测试 ===");
        
        long totalStartTime = System.currentTimeMillis();
        AtomicLong totalRecords = new AtomicLong(0);
        AtomicLong totalErrors = new AtomicLong(0);
        
        // 测试多个数据集
        String[] datasets = {"NOAA", "Stock", "IoT", "Network"};
        
        for (String dataset : datasets) {
            System.out.println("\n--- 测试数据集: " + dataset + " ---");
            
            try {
                long datasetStartTime = System.currentTimeMillis();
                AtomicLong datasetRecords = new AtomicLong(0);
                AtomicLong datasetErrors = new AtomicLong(0);
                
                // 根据数据集类型选择不同的测试方法
                switch (dataset) {
                    case "NOAA":
                        testNOAADataset(datasetRecords, datasetErrors);
                        break;
                    case "Stock":
                        testStockDataset(datasetRecords, datasetErrors);
                        break;
                    case "IoT":
                        testIoTDataset(datasetRecords, datasetErrors);
                        break;
                    case "Network":
                        testNetworkDataset(datasetRecords, datasetErrors);
                        break;
                }
                
                long datasetEndTime = System.currentTimeMillis();
                long datasetDuration = datasetEndTime - datasetStartTime;
                double datasetRecordsPerSecond = (double) datasetRecords.get() / datasetDuration * 1000;
                
                System.out.println(dataset + " 数据集测试完成:");
                System.out.println("记录数: " + datasetRecords.get());
                System.out.println("错误数: " + datasetErrors.get());
                System.out.println("耗时: " + datasetDuration + "ms");
                System.out.println("速度: " + String.format("%.2f", datasetRecordsPerSecond) + " 记录/秒");
                
                totalRecords.addAndGet(datasetRecords.get());
                totalErrors.addAndGet(datasetErrors.get());
                
            } catch (Exception e) {
                System.err.println(dataset + " 数据集测试失败: " + e.getMessage());
                totalErrors.incrementAndGet();
            }
        }
        
        long totalEndTime = System.currentTimeMillis();
        long totalDuration = totalEndTime - totalStartTime;
        double totalRecordsPerSecond = (double) totalRecords.get() / totalDuration * 1000;
        
        System.out.println("\n=== 综合测试结果 ===");
        System.out.println("总记录数: " + totalRecords.get());
        System.out.println("总错误数: " + totalErrors.get());
        System.out.println("总耗时: " + totalDuration + "ms");
        System.out.println("平均速度: " + String.format("%.2f", totalRecordsPerSecond) + " 记录/秒");
        System.out.println("成功率: " + String.format("%.2f", (double) totalRecords.get() / (totalRecords.get() + totalErrors.get()) * 100) + "%");
    }
    
    private void testNOAADataset(AtomicLong records, AtomicLong errors) throws IOException, KWDBException {
        String filepath = downloader.downloadNOAADataset();
        List<DataPoint> dataPoints = downloader.parseNOAADataset(filepath, 5000);
        
        for (DataPoint dataPoint : dataPoints) {
            try {
                client.write(dataPoint);
                records.incrementAndGet();
            } catch (Exception e) {
                errors.incrementAndGet();
            }
        }
    }
    
    private void testStockDataset(AtomicLong records, AtomicLong errors) throws IOException, KWDBException {
        String filepath = downloader.downloadStockPriceDataset();
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 1000);
        
        for (DataPoint dataPoint : dataPoints) {
            try {
                client.write(dataPoint);
                records.incrementAndGet();
            } catch (Exception e) {
                errors.incrementAndGet();
            }
        }
    }
    
    private void testIoTDataset(AtomicLong records, AtomicLong errors) throws IOException, KWDBException {
        String filepath = downloader.downloadIoTSensorDataset();
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 2000);
        
        for (DataPoint dataPoint : dataPoints) {
            try {
                client.write(dataPoint);
                records.incrementAndGet();
            } catch (Exception e) {
                errors.incrementAndGet();
            }
        }
    }
    
    private void testNetworkDataset(AtomicLong records, AtomicLong errors) throws IOException, KWDBException {
        String filepath = downloader.downloadNetworkMonitoringDataset();
        if (filepath.endsWith(".gz")) {
            filepath = downloader.extractGZFile(filepath);
        }
        List<DataPoint> dataPoints = downloader.parseCSVDataset(filepath, 3000);
        
        for (DataPoint dataPoint : dataPoints) {
            try {
                client.write(dataPoint);
                records.incrementAndGet();
            } catch (Exception e) {
                errors.incrementAndGet();
            }
        }
    }
}
